import tritonclient.grpc as grpcclient
from tritonclient import utils
from tritonclient.utils import InferenceServerException
import tritonclient.utils.cuda_shared_memory as cudashm
import sys
import json

triton_client = grpcclient.InferenceServerClient(
    url="0.0.0.0:8001",
    #url  = "172.17.0.1:8001",
    verbose=False,
    ssl=False,
    root_certificates=None,
    private_key=None,
    certificate_chain=None)

if not triton_client.is_server_live():
    print("FAILED : is_server_live")
    sys.exit(1)

if not triton_client.is_server_ready():
    print("FAILED : is_server_ready")
    sys.exit(1)

print (triton_client.get_cuda_shared_memory_status())
info = triton_client.get_cuda_shared_memory_status()
json_str=str(info).split("regions")[1:]
shm_num = 1
shm_region = []
print("现在正在使用的内存交换区域有")
for item in json_str:
    r_item = item.split('\n')[1][8:-1]
    print(str(shm_num),":",r_item)
    shm_num +=1
    shm_region.append(r_item)

a = -1
while(1):
    a = input("请输入想删除的数字序号(eg：1 2 3 4,空格区分)，输入0退出:")
    if(a =='0'):break

    del_list = list(map(int,a.split(' ')))
    for i in del_list:
        triton_client.unregister_cuda_shared_memory(shm_region[i-1])
    print("剩下的内存交换区域有:")
    info = triton_client.get_cuda_shared_memory_status()
    json_str=str(info).split("regions")[1:]
    shm_num = 1
    shm_region = []
    for item in json_str:
        r_item = item.split('\n')[1][8:-1]
        print(str(shm_num),":",r_item)
        shm_num +=1
        shm_region.append(r_item)
